import re
from typing import List, Dict, Optional
import tiktoken


class ContextManager:
    """对话上下文管理器"""
    
    def __init__(self, max_tokens: int = 4000, max_messages: int = 20, use_cache: bool = True):
        """
        初始化上下文管理器
        
        Args:
            max_tokens: 最大token限制
            max_messages: 最大消息数量限制
            use_cache: 是否使用上下文缓存
        """
        self.max_tokens = max_tokens
        self.max_messages = max_messages
        self.use_cache = use_cache
        self.encoding = None
        
        try:
            # 尝试加载tiktoken编码器
            self.encoding = tiktoken.get_encoding("cl100k_base")
        except:
            # 如果tiktoken不可用，使用简单的字符计数
            pass
    
    def count_tokens(self, text: str) -> int:
        """计算文本的token数量"""
        if self.encoding:
            return len(self.encoding.encode(text))
        else:
            # 简单估算：1个token ≈ 4个字符
            return len(text) // 4
    
    def count_message_tokens(self, message: Dict[str, str]) -> int:
        """计算单条消息的token数量"""
        role = message.get('role', '')
        content = message.get('content', '')
        thinking = message.get('thinking', '')
        
        # 计算基础token
        tokens = self.count_tokens(f"{role}: {content}")
        
        # 如果有思考过程，也计算在内
        if thinking:
            tokens += self.count_tokens(f"thinking: {thinking}")
            
        return tokens
    
    def count_conversation_tokens(self, messages: List[Dict[str, str]]) -> int:
        """计算整个对话的token数量"""
        total_tokens = 0
        for message in messages:
            total_tokens += self.count_message_tokens(message)
        return total_tokens
    
    def truncate_conversation(self, messages: List[Dict[str, str]], 
                            system_prompt: str = "") -> List[Dict[str, str]]:
        """
        智能截断对话历史，确保不超过token限制
        
        Args:
            messages: 完整的对话历史
            system_prompt: 系统提示词（用于计算总token）
            
        Returns:
            截断后的对话历史
        """
        if not messages:
            return messages
            
        # 计算系统提示词的token
        system_tokens = self.count_tokens(system_prompt) if system_prompt else 0
        
        # 从最新消息开始处理，保留最重要的上下文
        truncated_messages = []
        current_tokens = system_tokens
        
        # 首先确保包含最近的用户-助手对话对
        recent_pairs = []
        temp_messages = messages.copy()
        
        # 从后向前遍历，收集完整的对话对
        while temp_messages:
            msg = temp_messages.pop()
            if msg.get('role') == 'assistant':
                # 找到助手回复，向前找对应的用户消息
                assistant_msg = msg
                user_msg = None
                
                # 向前查找用户消息
                for prev_msg in reversed(temp_messages):
                    if prev_msg.get('role') == 'user':
                        user_msg = prev_msg
                        temp_messages.remove(prev_msg)
                        break
                
                if user_msg:
                    recent_pairs.insert(0, (user_msg, assistant_msg))
        
        # 添加最近的对话对（不超过最大限制）
        for user_msg, assistant_msg in recent_pairs[-self.max_messages//2:]:
            user_tokens = self.count_message_tokens(user_msg)
            assistant_tokens = self.count_message_tokens(assistant_msg)
            
            if current_tokens + user_tokens + assistant_tokens <= self.max_tokens:
                truncated_messages.extend([user_msg, assistant_msg])
                current_tokens += user_tokens + assistant_tokens
            else:
                break
        
        # 如果还有空间，添加一些系统消息
        remaining_tokens = self.max_tokens - current_tokens
        if remaining_tokens > 100:  # 至少100个token的空间
            for msg in reversed(messages):
                if msg.get('role') == 'system' and msg not in truncated_messages:
                    msg_tokens = self.count_message_tokens(msg)
                    if msg_tokens <= remaining_tokens:
                        truncated_messages.insert(0, msg)
                        remaining_tokens -= msg_tokens
        
        return truncated_messages
    
    def optimize_context(self, messages: List[Dict[str, str]], 
                       system_prompt: str = "", 
                       current_query: str = "") -> List[Dict[str, str]]:
        """
        优化对话上下文，包括智能截断和摘要
        
        Args:
            messages: 完整的对话历史
            system_prompt: 系统提示词
            current_query: 当前用户查询
            
        Returns:
            优化后的消息列表（包含系统提示词）
        """
        # 构建基础消息列表
        optimized_messages = []
        
        # 添加系统提示词
        if system_prompt:
            optimized_messages.append({"role": "system", "content": system_prompt})
        
        # 智能截断历史消息
        truncated_history = self.truncate_conversation(messages, system_prompt)
        optimized_messages.extend(truncated_history)
        
        # 如果历史消息被截断，添加摘要提示
        if len(truncated_history) < len(messages):
            summary_msg = self._create_summary_message(messages, truncated_history)
            if summary_msg:
                optimized_messages.append(summary_msg)
        
        return optimized_messages
    
    def _create_summary_message(self, full_history: List[Dict[str, str]], 
                              kept_history: List[Dict[str, str]]) -> Optional[Dict[str, str]]:
        """创建历史摘要消息"""
        # 找出被移除的消息
        removed_messages = [msg for msg in full_history if msg not in kept_history]
        
        if not removed_messages:
            return None
        
        # 提取被移除消息的关键信息
        summary_content = "以下是之前对话的摘要（由于长度限制，部分历史消息已被省略）：\n\n"
        
        user_count = sum(1 for msg in removed_messages if msg.get('role') == 'user')
        assistant_count = sum(1 for msg in removed_messages if msg.get('role') == 'assistant')
        
        summary_content += f"- 省略了 {user_count} 条用户消息和 {assistant_count} 条AI回复\n"
        
        # 添加一些关键的用户消息
        user_messages = [msg for msg in removed_messages if msg.get('role') == 'user']
        if user_messages:
            summary_content += "- 涉及的主要话题包括："
            topics = set()
            for msg in user_messages[:3]:  # 取前3条用户消息
                content = msg.get('content', '')
                # 提取关键词（简单实现）
                words = re.findall(r'\b\w{4,}\b', content.lower())
                topics.update(words[:3])  # 每条消息取前3个关键词
            
            if topics:
                summary_content += ", ".join(list(topics)[:5])  # 最多显示5个关键词
            summary_content += "\n"
        
        summary_content += "\n请基于当前上下文继续对话。"
        
        return {"role": "system", "content": summary_content}
    
    def get_recent_context(self, messages: List[Dict[str, str]], 
                         look_back: int = 3) -> List[Dict[str, str]]:
        """
        获取最近的对话上下文
        
        Args:
            messages: 完整的对话历史
            look_back: 回溯的消息对数（用户-助手对）
            
        Returns:
            最近的对话上下文
        """
        if not messages:
            return []
            
        recent_context = []
        pair_count = 0
        
        # 从后向前遍历，收集完整的对话对
        temp_messages = messages.copy()
        while temp_messages and pair_count < look_back:
            msg = temp_messages.pop()
            if msg.get('role') == 'assistant':
                # 找到助手回复，向前找对应的用户消息
                assistant_msg = msg
                user_msg = None
                
                # 向前查找用户消息
                for prev_msg in reversed(temp_messages):
                    if prev_msg.get('role') == 'user':
                        user_msg = prev_msg
                        temp_messages.remove(prev_msg)
                        break
                
                if user_msg:
                    recent_context.extend([user_msg, assistant_msg])
                    pair_count += 1
        
        # 反转顺序，保持时间顺序
        return list(reversed(recent_context))

    def set_max_tokens(self, max_tokens: int):
        """设置最大token限制"""
        self.max_tokens = max_tokens

    def set_max_messages(self, max_messages: int):
        """设置最大消息数量限制"""
        self.max_messages = max_messages

    def set_use_cache(self, use_cache: bool):
        """设置是否使用上下文缓存"""
        self.use_cache = use_cache

    def get_settings(self) -> Dict[str, any]:
        """获取当前设置"""
        return {
            "max_tokens": self.max_tokens,
            "max_messages": self.max_messages,
            "use_cache": self.use_cache
        }


# 创建全局上下文管理器实例
context_manager = ContextManager()